import pandas as pd
import matplotlib.pyplot as plt
import sys
import os
import re

def sanitize_filename(name):
    """清理文件名，去掉非法字符"""
    name = re.sub(r'[\\/*?:"<>|]', "_", name)
    return name

def plot_all_ops(csv_file, batch_sizes_to_plot, output_dir):
    df = pd.read_csv(csv_file)

    os.makedirs(output_dir, exist_ok=True)

    all_ops = sorted(df['op_name'].dropna().unique())

    for op_name in all_ops:
        df_op = df[df['op_name'] == op_name]

        for batch_size in batch_sizes_to_plot:
            df_batch = df_op[df_op['batch_size'] == batch_size]

            if df_batch.empty:
                continue

            # 取出所有出现过的 threads_host，按数值排序
            all_threads = sorted(df_batch['threads_host'].unique())
            thread_labels = [str(t) for t in all_threads]  # 用字符串作为分类标签

            fig, ax = plt.subplots(figsize=(10, 6))

            for device_type in ['dpu', 'dpa', 'cpu']:
                df_device = df_batch[df_batch['device_type'] == device_type]
                if df_device.empty:
                    continue

                # 重新对齐，确保每个 threads_host 都有一行
                df_device = df_device.set_index('threads_host').reindex(all_threads).reset_index()

                x = range(len(all_threads))  # x 变成 0,1,2,3...，均匀分布
                y = df_device['p50_ms']

                # 处理超过100ms的点
                y_plot = []
                for yi in y:
                    if pd.isna(yi):
                        y_plot.append(None)
                    elif yi > 100:
                        y_plot.append(110)  # 超过100ms的点，画到110
                    else:
                        y_plot.append(yi)

                ax.plot(x, y_plot, marker='o', label=device_type.upper())

                # ★★★ 每个点都标注数值 ★★★
                for xi, yi, yi_plot in zip(x, y, y_plot):
                    if pd.notna(yi):
                        display_text = f'{yi:.1f}ms'
                        # 如果是超过100ms的点，标注位置稍微高一点
                        offset = 5 if yi > 100 else 3
                        ax.text(xi, yi_plot + offset, display_text, ha='center', va='bottom', fontsize=8)

            # 设置横坐标为分类型
            ax.set_xticks(range(len(all_threads)))
            ax.set_xticklabels(thread_labels)

            ax.set_xlabel('Threads Host')
            ax.set_ylabel('P50 Execute Time (ms)')
            ax.set_title(f'{op_name} (Batch Size={batch_size})')

            ax.grid(True, linestyle='--', linewidth=0.5)

            # 设置纵坐标范围
            ax.set_ylim(0, 130)

            ax.legend()

            sanitized_op_name = sanitize_filename(op_name)
            output_base = os.path.join(output_dir, f'{sanitized_op_name}_batch{batch_size}')
            plt.savefig(output_base + '.png')
            # plt.savefig(output_base + '.pdf')  # 如果需要保存 PDF，可以取消注释
            plt.close()
            print(f"Saved plot to {output_base}.png")

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python3 plot_all_ops.py <summary_csv> <output_dir> <batch_size1> [<batch_size2> ...]")
        sys.exit(1)

    csv_file = sys.argv[1]
    output_dir = sys.argv[2]
    batch_sizes = list(map(int, sys.argv[3:]))

    plot_all_ops(csv_file, batch_sizes, output_dir)
